%% This file is called to compute comparative static results; 

disp('-------------> Model 6 (CWZ Model), idiosyncratic shocks + imperfect reallocation + imperfect credit + money <---------')
%%%% It needs "par" as parameters
%%%% Outputs results stored in variables named similar to Inv_model_6 below

%% Preparation
switch variable_x
    case '\alpha'
        l_x         = length(ALPHA_grid);
        x_grid      = ALPHA_grid;
 
    case '\tau_k'
        l_x         = length(tau_k_grid);
        x_grid      = tau_k_grid;
    
    case '\iota'
        l_x         = 30;           % total number of grids 
        l_mon       = 30;           % number of grids for monetary equilibrium
        l_non_mon   = l_x - l_mon;  % number of grids for non-monetary equilibrium
        iota_grid   = zeros(l_x, length(THETA_grid));
        x_grid      = iota_grid;      
end


%% Initilization
para_XI              = zeros(length(THETA_grid),1); %%%% used for welfare analysis, as the parameters could change each time with a different theta.
para_XI_IM           = para_XI;

Inv_model_6          = zeros(l_x, length(THETA_grid)); 
Threshold_model_6    = Inv_model_6; 
Prod_model_6         = Inv_model_6;  
Prod_std_model_6     = Inv_model_6;
Output_model_6       = Inv_model_6; 
ALPHA_0_model_6      = Inv_model_6; 
IM_model_6           = Inv_model_6;
Hours_model_6        = Inv_model_6; 
R_ratio_model_6      = Inv_model_6; 
P_share_model_6      = Inv_model_6;
Reallocation_model_6 = Inv_model_6; 
debt_Y_model_6       = Inv_model_6; 
debt_model_6         = Inv_model_6;
B_model_6            = Inv_model_6; 
Aqc_P_model_6        = Inv_model_6; 
Aqc_model_6          = Inv_model_6; 
Sppe_model_6         = Inv_model_6; 
Sppe_P_model_6       = Inv_model_6; 
Sec_mkt_price_model_6= Inv_model_6;
Util_model_6         = Inv_model_6; 
Welfare_model_6      = Inv_model_6; 
Consumption_model_6  = Inv_model_6;
CM_C_model_6         = Inv_model_6; 
Z_Y_model_6          = Inv_model_6; 
Z_K_model_6          = Inv_model_6;
monetary_eqm_model_6 = Inv_model_6;
variable_x_threshold = max(x_grid) .* ones(1, length(THETA_grid)) * 100;



%% Running comparative statics for different theta

for ii = 1: length(THETA_grid)
    par         = par_ini;
    par.THETA   = THETA_grid(ii);    
    fun         = define_function_fun(par); %%%% note: changing iota does not change those functions
    
    disp(' ')
    fprintf('----- The current theta is %.4f ---- \n', par.THETA)
  
    if strcmp(variable_x,'\iota') % if the computation is based on different iota
        disp('I am doing a new calibration')
        do_calibration_main_model; %%%% recalibrate each time of a new THETA; comment it out if you don't want this step
        
        para_XI(ii)              = par.XI; %%%% recalibration means that XI changes
        para_XI_IM(ii)           = par.XI_IM;   
        
        %%%% First, try the non-monetary eqm
        perfect_reallocation     = 0; 
        cali_step                = 0;
        x_non                    = fsolve(@(x)eqm_ss_iid_cases_2_to_5(x, par, fun, perfect_reallocation, cali_step), [0.1, 0], options);     
        B_non                    = x_non(1); 
        ALPHA_0_non              = x_non(2);

        [fval, eqm_non] = eqm_ss_iid_cases_2_to_5(x_non, par, fun, perfect_reallocation, 0);
        
        variable_x_threshold(ii) = eqm_non.threshold_iota;

        iota_grid(:, ii)         = [linspace(0.0001, eqm_non.threshold_iota-0.0005, l_mon), linspace(eqm_non.threshold_iota + 0.0001, 0.2, l_non_mon)]';
        
        %%%% Second, try the monetary eqm
        x0 = [par.EPS_L_BOUND * 4 + 1.2, B_non, ALPHA_0_non]; %%%% as initial guess for solving monetary equilibrium     
        parfor jj = 1 : l_x
            par_temp       = par;
            par_temp.iota  = iota_grid(jj, ii);
            
            if eqm_non.threshold_iota > par_temp.iota %%%% monetary equilibrium's outcome
                x       = fsolve(@(x)eqm_ss_iid_case_6(x, par_temp, fun, 1, 0), x0, options);
                L       = x(1);
                B       = x(2);
                ALPHA_0 = x(3);
                monetary_eqm_model_6(jj, ii) = 1;
                [fval, eqm] = eqm_ss_iid_case_6(x, par_temp, fun, 1, 0);           
            else %%%% non-monetary equilibrium's outcome
                eqm = eqm_non;
            end
            
            eqm.IM                      = average_log_normal(eqm.surplus_net, par.c_mu, par.c_sig) / par.XI_IM;      
            Consumption_model_6(jj,ii)  = eqm.c_star + eqm.w * eqm.IM * par_temp.ZETA;
            CM_C_model_6(jj,ii)         = eqm.c_star;            
            Inv_model_6(jj,ii)          = eqm.K * par_temp.DELTA;
            Hours_model_6(jj,ii)        = eqm.H;
            
            Output_model_6(jj,ii)       = eqm.Y + eqm.w * eqm.IM * par.ZETA; 
            IM_model_6(jj, ii)          = eqm.IM;
            
            Prod_model_6(jj,ii)         = eqm.K_bar_K / eqm_calibrated.K_bar_K; %%%% relative to calibrated level of productivity
            Prod_std_model_6(jj,ii)     = eqm.prod_std;
   
            R_ratio_model_6(jj,ii)      = (eqm.E_a_K + eqm.E_p_K) / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            P_share_model_6(jj,ii)      = eqm.E_p_K / (eqm.E_a_K + eqm.E_p_K);
            Aqc_model_6(jj,ii)          = eqm.E_a_K  / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            Aqc_P_model_6(jj,ii)        = eqm.prob_A;
            Sppe_model_6(jj,ii)         = eqm.E_p_K  / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            Sppe_P_model_6(jj,ii)       = eqm.prob_P;
            Util_model_6(jj,ii)         = log(eqm.c_star) - par_temp.XI * eqm.H - par_temp.XI_IM * eqm.IM;
            Welfare_model_6(jj,ii)      = eqm.c_star / eqm_calibrated.c_star * exp(par_temp.XI * (eqm_calibrated.H - eqm.H) + par_temp.XI_IM * (eqm_calibrated.IM - eqm.IM)) - 1; % use calibrated c_star and H
            
            Sec_mkt_price_model_6(jj,ii) = eqm.average_2nd_price;
            Z_Y_model_6(jj,ii)          = eqm.ALPHA_0 * eqm.Z / eqm.Y;
            Z_K_model_6(jj,ii)          = eqm.ALPHA_0 * eqm.Z / eqm.K;
            debt_Y_model_6(jj,ii)       = eqm.ALPHA_0 * par_temp.ALPHA * eqm.debt / eqm.Y; % debt occurs only if you enter the DM and meet someone
            debt_model_6(jj,ii)         = eqm.ALPHA_0 * par_temp.ALPHA * eqm.debt; % debt occurs only if you enter the DM and meet someone        
            ALPHA_0_model_6(jj,ii)      = eqm.ALPHA_0;
            Threshold_model_6(jj,ii)    = eqm_non.threshold_iota;
        end

    else  % not on different \iota
        x0 = [par.EPS_L_BOUND * 4 + 1, 0.08, 0.05];
        for jj = 1 : l_x
            par_temp = par;
 
            switch variable_x
                case '\alpha'
                    par_temp.ALPHA  = ALPHA_grid(jj); 
                case '\tau_k'
                    par_temp.tau_k  = tau_k_grid(jj);
            end
            
            fun = define_function_fun(par_temp);
            
            perfect_reallocation     = 0; 
            cali_step                = 0;
            x_non                    = fsolve(@(x)eqm_ss_iid_cases_2_to_5(x, par_temp, fun, perfect_reallocation, cali_step), [0.08, 0.01], options);
            B_non                    = x_non(1);
            ALPHA_0_non              = x_non(2);
            
            [fval, eqm_non] = eqm_ss_iid_cases_2_to_5(x_non, par_temp, fun, perfect_reallocation, cali_step);
            if eqm_non.threshold_iota > par_temp.iota + 0.0001 %%%% monetary equilibrium then can exist
                x       = fsolve(@(x)eqm_ss_iid_case_6(x, par_temp, fun, 1, 0), x0, options);
                x0      = x;
                L       = x(1);
                B       = x(2);
                ALPHA_0 = x(3);
                monetary_eqm_model_6(jj,ii) = 1;              
                [fval, eqm] = eqm_ss_iid_case_6(x, par_temp, fun, 1, 0);  
            else
                eqm = eqm_non;
            end

            eqm.IM                      = average_log_normal(eqm.surplus_net, par.c_mu, par.c_sig) / par.XI_IM;      
            Consumption_model_6(jj,ii)  = eqm.c_star + eqm.w * eqm.IM * par_temp.ZETA;
            CM_C_model_6(jj,ii)         = eqm.c_star;            
            Inv_model_6(jj,ii)          = eqm.K * par_temp.DELTA ;
            Hours_model_6(jj,ii)        = eqm.H;
            
            Output_model_6(jj,ii)       = eqm.Y + eqm.w * eqm.IM * par.ZETA; %%%% the normalization is as if
            IM_model_6(jj, ii)          = eqm.IM;
            
            Prod_model_6(jj,ii)         = eqm.K_bar_K / eqm_calibrated.K_bar_K; %%%% relative to calibrated level of productivity
            Prod_std_model_6(jj,ii)     = eqm.prod_std;
   
            R_ratio_model_6(jj,ii)      = (eqm.E_a_K + eqm.E_p_K) / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            P_share_model_6(jj,ii)      = eqm.E_p_K / (eqm.E_a_K + eqm.E_p_K);
            Aqc_model_6(jj,ii)          = eqm.E_a_K  / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            Aqc_P_model_6(jj,ii)        = eqm.prob_A;
            Sppe_model_6(jj,ii)         = eqm.E_p_K  / (eqm.E_a_K + eqm.E_p_K + par_temp.DELTA);
            Sppe_P_model_6(jj,ii)       = eqm.prob_P;
            Util_model_6(jj,ii)         = log(eqm.c_star) - par_temp.XI * eqm.H - par_temp.XI_IM * eqm.IM;
            Welfare_model_6(jj,ii)      = eqm.c_star / eqm_calibrated.c_star * exp(par_temp.XI * (eqm_calibrated.H - eqm.H) + par_temp.XI_IM * (eqm_calibrated.IM - eqm.IM)) - 1; % use calibrated c_star and H
            
            Sec_mkt_price_model_6(jj,ii)= eqm.average_2nd_price;
            Z_Y_model_6(jj,ii)          = eqm.ALPHA_0 * eqm.Z / eqm.Y;
            Z_K_model_6(jj,ii)          = eqm.ALPHA_0 * eqm.Z / eqm.K;
            debt_Y_model_6(jj,ii)       = eqm.ALPHA_0 * par_temp.ALPHA * eqm.debt / eqm.Y; %%%% debt occurs only if you enter the DM and meet someone
            debt_model_6(jj,ii)         = eqm.ALPHA_0 * par_temp.ALPHA * eqm.debt; %%%% debt occurs only if you enter the DM and meet someone
            ALPHA_0_model_6(jj,ii)      = eqm.ALPHA_0;
            Threshold_model_6(jj,ii)    = eqm_non.threshold_iota;
        end

        idx = find(monetary_eqm_model_6(:,ii) == 0, 1, 'last') +1; %%%% unlike iota, higher ALPHA can generate monetary eqm
        if isempty(idx)
            continue
        end
        variable_x_threshold(ii) = x_grid(idx);
    end
end

%% plotting
x_grid_temp = x_grid;

if strcmp(variable_x,'\iota') || strcmp(variable_x, 'inflation')
    x_grid = iota_grid;

    %%%% inflation or interest rate
    adj = - inflation_ * (1 / par.BETA - 1);
    range = [0, 0.18];
    if inflation_ ==1
        variable_x = '\pi';
        range = [0, 15] / 100;
    end

else
    adj = 0;
    x_grid = repmat(x_grid_temp', 1, length(THETA_grid)); 
end

if plotting_ == 1
    %%%%  plotting long-run effects of inflation
    figure()  

    subplot(3,4,1)
    hold on
    for ii = 1:length(THETA_grid)      
        id1 = x_grid(:,ii) < variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Output_model_6(id1,ii), 'LineWidth', 2);    
        h2 = plot(x_grid(id2, ii) + adj, Output_model_6(id2,ii), 'LineWidth', 2); 
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Output')
    xlabel(variable_x)
    xlim(range)
    grid on

    subplot(3,4,2)
    hold on
    for ii = 1:length(THETA_grid)
        investment_specific_adjustment = 0;
        id1 = x_grid(:,ii) < variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Inv_model_6(id1,ii), 'LineWidth', 2);    
        h2 = plot(x_grid(id2, ii) + adj, Inv_model_6(id2,ii), 'LineWidth', 2); 
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
        %%%%note: we want to show investment different movement with inflation;
        %%%%so we do not adjust by adding adj
    end
    title('Investment')
    xlabel(variable_x)
    xlim(range)
    grid on

    subplot(3,4,3)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Consumption_model_6(id1,ii), 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Consumption_model_6(id2,ii), 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Consumption')
    xlabel(variable_x)
    xlim(range)
    grid on

    subplot(3,4,4)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Hours_model_6(id1,ii) * 100, 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Hours_model_6(id2,ii) * 100, 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Employment')
    xlabel(variable_x)
    ylabel('%')
    xlim(range)
    grid on

    subplot(3,4,5)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, R_ratio_model_6(id1,ii), 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, R_ratio_model_6(id2,ii), 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('R share')
    xlabel(variable_x)
    xlim(range)
    ylim([0 1])
    grid on

    subplot(3,4,6)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, P_share_model_6(id1,ii), 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, P_share_model_6(id2,ii), 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});

        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('P share')
    xlabel(variable_x)
    xlim(range)
    grid on    

    subplot(3,4,7)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Aqc_P_model_6(id1,ii), 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Aqc_P_model_6(id2,ii), 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}}); 
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end 
    title('Prob of full sale')
    xlabel(variable_x)
    xlim(range)
    grid on

    subplot(3,4,8)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Sec_mkt_price_model_6(id1,ii), 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Sec_mkt_price_model_6(id2,ii), 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}}); 
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('DM Price')
    xlabel(variable_x)
    xlim(range)
    grid on

    subplot(3,4,9)
    hold on
    for ii = 1:length(THETA_grid)    
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Prod_model_6(id1,ii)  * 100, 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Prod_model_6(id2,ii)  * 100, 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});   

        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Productivity')
    xlabel(variable_x)
    ylabel('%')
    xlim(range)
    grid on

    subplot(3,4,10)
    hold on
    for ii = 1:length(THETA_grid) 
        id1 = x_grid(:,ii) < variable_x_threshold(ii);
        id2 = ~id1;
        fun_temp = @(Xn) interp1(x_grid(id1, ii), Util_model_6(id1, ii), Xn);
        if strcmp(variable_x,'\iota') || strcmp(variable_x,'\pi')%%%% when varying inflation, zero inflation is the benchmark!!
            W_0 = fun_temp(1 / par.BETA - 1); % 
            h1 = plot(x_grid(id1, ii) + adj, exp(W_0 + para_XI(ii) * Hours_model_6(id1, ii) + para_XI_IM(ii) * IM_model_6(id1, ii)) ./ CM_C_model_6(id1, ii) * 100 - 100, 'LineWidth', 2);
            h2 = plot(x_grid(id2, ii) + adj, exp(W_0 + para_XI(ii) * Hours_model_6(id2, ii) + para_XI_IM(ii) * IM_model_6(id2, ii)) ./ CM_C_model_6(id2, ii) * 100 - 100, 'LineWidth', 2);
        else
            h1 = plot(x_grid(id1, ii), -Welfare_model_6(id1, ii) * 100, 'LineWidth', 2);
            h2 = plot(x_grid(id2, ii), -Welfare_model_6(id2, ii) * 100, 'LineWidth', 2);
        end

        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});   
        xline(variable_x_threshold(ii) + adj, '-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Welfare cost (equiv C)')
    xlabel(variable_x)
    ylabel('%')
    xlim(range)
    grid on
    
    subplot(3,4,11)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1 = plot(x_grid(id1, ii) + adj, Z_Y_model_6(id1,ii) * 100, 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, Z_Y_model_6(id2,ii) * 100, 'LineWidth', 2);
        set(h1, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Cash / output')
    xlabel(variable_x)
    ylabel('%')
    xlim(range)
    ylim([0, 100])
    grid on

    subplot(3,4,12)
    hold on
    for ii = 1:length(THETA_grid)
        id1 = x_grid(:,ii)<variable_x_threshold(ii);
        id2 = ~id1;
        h1(ii) = plot(x_grid(id1, ii) + adj, debt_Y_model_6(id1,ii) * 100, 'LineWidth', 2);
        h2 = plot(x_grid(id2, ii) + adj, debt_Y_model_6(id2,ii) * 100, 'LineWidth', 2);
        set(h1(ii), {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        set(h2, {'color'}, num2cell(colors(ii,:), 2), {'LineStyle'}, {line_style{ii}});
        xline(variable_x_threshold(ii) + adj,'-.','color',colors(ii,:),'LineWidth', 1);
    end
    title('Debt / output')
    xlabel(variable_x)
    ylabel('%')
    xlim(range)
    grid on

    for i = 1:length(THETA_grid)
        entry_legend{i} =  ['\theta = ' THETA_grid_legend{i}];
    end
    legend(h1,entry_legend, 'fontsize',14, 'location','Northeast','Orientation','horizontal');
end
